import json
from enum import Enum
from typing import List, Optional, Any, Dict, Union

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore, VectorStoreRetriever
from loguru import logger


class VectorType(str, Enum):
    ELASTICSEARCH = "Elasticsearch"
    MILVUS = "Milvus"
    OpenSearch = "OpenSearch"


class VectorStoreClient:
    """向量数据库操作客户端"""

    def __init__(self,
                 embeddings: Embeddings,
                 collection_name: str,
                 connect_config: dict):
        self._embeddings = embeddings
        self._collection_name = collection_name
        self._connect = connect_config
        self._vector_type = connect_config["type"]
        self._vector_store: Optional[VectorStore] = None
        logger.info("连接host: " + str(self._connect["host"]))

    def get_vector_store(self) -> VectorStore:
        return self._vector_store

    def get_retriever(self, **kwargs: Any) -> VectorStoreRetriever:
        """从向量存储构建检索器"""
        return self._vector_store.as_retriever(**kwargs)

    def get_vector_type(self) -> VectorType:
        return self._vector_type

    def get_embeddings(self) -> Embeddings:
        return self._embeddings

    def get_connect(self) -> Optional[dict[str, Any]]:
        return self._connect

    def get_collection_name(self) -> str:
        return self._collection_name

    def add_documents(self, docs: List[Document], **kwargs: Any):
        pass

    def delete_documents(self,
                         body: Any) -> Optional[bool]:
        pass

    def query_documents(self, body: Any) -> List[Document]:
        pass

    def create_collection(self, mappings: Dict[str, str]):
        pass
